# -*- coding: utf-8 -*-

from WSModule import WSModule
from WSThread import WSThread
from WSJob import WSJob
from WSCounter import WSCounter
from WSResult import WSResult
from WSOption import WSOption
from DnsBruteThread import DnsBruteThread
from DnsBruteResult import DnsBruteResult
from DnsBruteSrvs import DnsBruteSrvs
from DnsBruteJob import DnsBruteJob
from Registry import Registry
from DictOfMask import DictOfMask

import queue
import dns.query
import dns.message
import dns.resolver
import re
import os
import datetime
import time
# TODO - тест на все ошибки параметров
# TODO - speed TEST DNS
# В примечания - OpenDNS-сервера направляют на себя несуществующие домены
# А надо ли парсить вывод? Может там есть какие-то св-ва

#TODO параметр вывода результатов на экран
#TODO счётчики везде

class DnsBruteException(Exception): pass
# TODO привести в порядок все имена параметров и их описания
# TODO везде счётчики
# TODO вывод общего кол-ва очереди при счётчиках
class DnsBruteMask(WSModule):
    log_path = os.getcwd() + '/results/DNSBrute-result' # Заменить на аналог из реестра
    options = {"threads": WSOption("threads", "count thread", 10, False, ['-t', '--threads']),
               "mask": WSOption("mask", "Mask for brute", "", True, ['-m', '--mask']),
               "domain": WSOption("domain", "target domain", "vega-int.ru", True, ['-n', '--host']),
               "msymbol": WSOption("msymbol", "Symbol of mask position in target domain name (default №)", "№", False, ['--msymbol']),
               "dns_srv": WSOption("dns_srv", "work DNS-server", "8.8.8.8", False, ['-d', '--dns'])}

    dns_srv = None

    def run(self):
        if self.options['domain'].value.find(self.options['msymbol'].value) == -1:
            exit(print("Symbol of object position ({0}) not found in URL ({1}) ".format(self.options['msymbol'].value,
                                                                                        self.options['domain'].value)))
        self.dns_srv = DnsBruteSrvs(self.options['dns_srv'].value)

        self.job = DnsBruteJob()

        dom = DictOfMask(self.options['mask'].value)
        dict_for_brute = dom.dict()

        for word in dict_for_brute:
            self.job.put(word)

        self.result = DnsBruteResult()

        counter = WSCounter(5, 150)

        workers = []
        for i in range(self.options['threads'].value):
            workers.append(DnsBruteThread(self.job, self.result, self.options['domain'].value, self.options['msymbol'].value, self.dns_srv, counter))
        self.kernel.createThreads(workers)

    def help(self):
        return "Module DNSBrute ver 1.0"

    def finished(self):
        return self.kernel.finished()

